from PIL import Image, ImageDraw, ImageFont

from PIL import Image, ImageDraw, ImageFont

def draw_label_info_on_image(
    image_path, conf, label, label_map, output_path,
    font_path='./SarasaMonoSC-Regular.ttf', font_size=1.5
):
    # 读取图片并上采样到448x448
    image = Image.open(image_path).convert("RGB")
    image = image.resize((448, 448), Image.BICUBIC)

    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype(font_path, int(font_size * 10))  # 你可以自己微调下这个公式

    conf_text = f"置信度: {conf:.2f}"
    label_text = f"公开标签: {label_map.get(str(label), str(label))}"
    x, y = 5, 5
    line_spacing = int(font_size * 18)

    draw.text((x, y), conf_text, font=font, fill=(255, 255, 255))
    draw.text((x, y + line_spacing), label_text, font=font, fill=(255, 255, 255))
    image.save(output_path)

if __name__ == "__main__":
    
    # 生成一张黑图
    test_img = Image.new('RGB', (224, 224), (0, 0, 0))
    test_img.save('test_black.png')
    # 试试函数
    label_map = {str(i): f"标签{i}" for i in range(30)}
    draw_label_info_on_image('test_black.png', 0.23456, 12, label_map,  'test_black_labeled.png')
